import os
from argparse import ArgumentParser

from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt

from tools.json_helper import json_helper
from utils.preprocess import bilatera_blur


def open_img_np(
    src_file, *, 
    max_ : np.ndarray = np.array([255., 255., 255.]),
    min_ : np.ndarray = np.array([0., 0., 0.]),  
    scale : int = 255, mean=False, dims=True, **kwargs
) -> np.ndarray:
    """
    open image as numpy ndarray

    Param
    -----
    max_ :  max of each channel to normalize, no normalize when sum is 0
            max channels > file channels
    min_ :  min of each channel to normalize
            min channels > file channels
    scale : change to range of [0 scale] after normalize
    mean :  return mean of channels
    dims :  keep dims == 3

    Note
    -----
    plt.imread read png as float with [0, 1]
    otherwise int 
    """
    
    img_np = plt.imread(src_file)

    if len(img_np.shape) == 3 and img_np.shape[2] < max_.shape[0]:
        max_ = max_[:img_np.shape[2]]
        min_ = min_[:img_np.shape[2]]
    elif len(img_np.shape) == 2 and max_.shape[0] > 1:
        max_ = max_[0]
        min_ = min_[0]

    if max_.sum() != 0:
        img_np = (img_np.clip(*(min_, max_)) - min_) / (max_ - min_)

    if mean and len(img_np.shape) == 3:
        img_np = img_np.mean(axis=2)

    if max_.sum() != 0:
        img_np = (img_np * scale).astype(np.uint8)

    if dims and len(img_np.shape) == 2:
        img_np = np.expand_dims(img_np, 2)

    return img_np


def plot_crn(img_dict, save_dir, *, 
    color=[255, 0, 0], save_kw='', modal='', pproc_func=None,
    read_func=open_img_np, wkt2fn=None
):
    '''
    根据城市对应的img2pt绘图

    Param:
    ----
    img_dict: img2dict字典
    save_dir: 图像保存的路径
    color: 特征点标记的颜色，shape (h, w, 3)
    save_kw: 保存图像时的前缀
    modal: 模态名('opt'/'sar')
    pproc_func: 预处理函数，为None则不处理
    '''

    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)


    for kwt, points in tqdm(img_dict.items()):

        fn = wkt2fn(kwt, modal=modal)
        background = read_func(fn)

        if len(background.shape) == 2:
            background = np.expand_dims(background, axis=2)
        if background.shape[2] == 1:
            background = np.concatenate(
                (background, background, background), axis=2
            )

        background = pproc_func(background) if pproc_func != None else background

        fig = plt.figure()
        plt.imshow(background)
        plt.axis('off')

        max_response = -1
        for pt_attri in points.values():
            
            cur_coord = pt_attri["xy"]
            kpt_coord = [
                cur_coord[0]-5, cur_coord[1]-5, 
                cur_coord[0]+5, cur_coord[1]+5
                ]
            annot = round(pt_attri["response"], 2)
            max_response = max(max_response, annot)

            plt.gca().add_patch(
                plt.Rectangle((kpt_coord[0], kpt_coord[1]), 
                    kpt_coord[2] - kpt_coord[0],
                    kpt_coord[3] - kpt_coord[1], fill=False,
                    edgecolor='r', linewidth=1)
                )
            plt.text(
                kpt_coord[0], kpt_coord[1]-1, str(annot), 
                # backgroundcolor = color, 
                color = (1, 1, 1),  bbox=dict(boxstyle='round,pad=0', fc=color, lw=0,  alpha=0.7)
            )

        showimg_name = os.path.basename(fn)
        showimg_name = os.path.splitext(showimg_name)[0] + ".png"
        showimg_name = save_kw + showimg_name if save_kw != '' else showimg_name
        
        plt.title(f"max : {max_response}")
        # plt.imsave(os.path.join(save_dir, showimg_name), background)
        fig.tight_layout()
        plt.savefig(os.path.join(save_dir, showimg_name))
        plt.close(fig)


def plot_cmd(wkt2fn=lambda x: x, read_func=open_img_np):

    parser = ArgumentParser()
    parser.add_argument(
        '-o', '--opt_file', help='path of opt files', default=""
    )
    parser.add_argument(
        '-s', '--sar_file', help='path of sar files', default=""
    )
    parser.add_argument(
        '-v', '--save_dir', help='folder to save plot image', default=""
    )
    parser.add_argument(
        '-d', '--dset_dir', help='folder of dataset', default=""
    )

    args = parser.parse_args()

    img2pt_file_list = [args.opt_file, args.sar_file]
    dset_dir = args.dset_dir
    save_dir = args.save_dir
    modal_list = ["opt", "sar"]

    img2pt_list = []

    def wkt2fullname(wkt: str, modal:str):
        fn = wkt2fn(wkt, modal=modal)
        fn = os.path.join(dset_dir, fn)

        return fn

    for fn in img2pt_file_list:
        j = json_helper(json_dir=fn)
        img2pt_list.append(j.dict_)

    for img2pt, m in zip(img2pt_list, modal_list):
        pproc_func = bilatera_blur if m == "sar" else None 

        plot_crn(
            img2pt, save_dir, color=[1., 0, 0], save_kw=m, 
            modal=m, pproc_func=pproc_func,
            read_func=read_func, wkt2fn=wkt2fullname
        )


if __name__ == "__main__1":
    # png = "E:/datasets/QXSLAB_SAROPT/sar_256_oc_0.2/4.png"
    png = "E:/datasets/QXSLAB_SAROPT/opt_256_oc_0.2/1.png"

    img = open_img_np(png, max_=np.array([0., 0., 0., 1., 1.]), mean=False, dims=False)

    print(img.shape, img.max())


if __name__ == "__main__":

    from functools import partial
    from tools.QXSLAB_helper import open_QXSLAB, wkt2fn_QXSLAB
    from tools.json_helper import json_helper
    from utils.preprocess import bilatera_blur

    img2pt_file_list = [
        "E:\workspace\SOMatch\json\QXSLAB_QingDao_harris\OPT_HARRIS0.1.json", 
        "E:\workspace\SOMatch\json\QXSLAB_QingDao_harris\SAR_HARRIS0.1.json"
    ]
    save_dir = "E:\datasets\client-data\QXSLAB_QingDao_o1e-1s1e-1"
    dset_dir = "E:\datasets\QXSLAB_SAROPT"
    modal_list = ["opt", "sar"]

    img2pt_list = []

    def wkt2fullname(wkt: str, modal:str):
        fn = wkt2fn_QXSLAB(wkt, modal=modal)
        fn = os.path.join(dset_dir, fn)

        return fn

    for fn in img2pt_file_list:
        j = json_helper(json_dir=fn)
        img2pt_list.append(j.dict_)

    for img2pt, m in zip(img2pt_list, modal_list):
        pproc_func = bilatera_blur if m == "sar" else None 

        plot_crn(
            img2pt, save_dir, color=[1., 0, 0], save_kw=m, 
            modal=m, pproc_func=pproc_func,
            read_func=partial(open_QXSLAB, scale=0, mean=True), 
            wkt2fn=wkt2fullname
        )

